import os
import torch
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from torch import nn
from torch.utils.data import Dataset, DataLoader, TensorDataset, Subset


class TimeseriesData(object):
    def __init__(self):
        self.root_path = './'
        self.data_path = 'AAPL.csv'
        self.scaler = MinMaxScaler(feature_range=(0, 1))

    def _read(self):
        return pd.read_csv(
            os.path.join(self.root_path, self.data_path),
            parse_dates=True, )

    def values(self):
        values = self._read()[['Open', 'High', 'Low', 'Close']].values
        return self.scaler.fit_transform(values)

    def labels(self):
        labels = self._read()['Close'].values.reshape(-1, 1)
        return self.scaler.fit_transform(labels)


class SupervisedTimeseriesData(TimeseriesData):
    def __init__(self, lag=10):
        super(SupervisedTimeseriesData, self).__init__()
        self.lag = lag
        self.timeseries = TimeseriesData()

    def values(self):
        x_list =[]
        for i in range(len(self.timeseries.values()) - self.lag):
            x_list.append(self.timeseries.values()[i:i + self.lag])
        return torch.tensor(np.array(x_list), dtype=torch.float)

    def labels(self):
        y_list = []
        for i in range(len(self.timeseries.labels()) - self.lag):
            y_list.append(self.timeseries.labels()[i + self.lag])
        return torch.tensor(np.array(y_list), dtype=torch.float)


class SupervisedTimeseriesDataset(Dataset):
    def __init__(self):
        super(SupervisedTimeseriesDataset, self).__init__()
        self.set = SupervisedTimeseriesData()

    def __len__(self):
        return len(self.set.values())

    def __getitem__(self, index):
        return self.set.values()[index], self.set.labels()[index]

    def train_set(self):
        dataset = TensorDataset(self.set.values(), self.set.labels())
        return Subset(dataset, indices=list(range(self.__len__()//5*4)))

    def val_set(self):
        dataset = TensorDataset(self.set.values(), self.set.labels())
        return Subset(dataset, indices=list(range(self.__len__()//5*4, self.__len__()//10*9)))

    def test_set(self):
        dataset = TensorDataset(self.set.values(), self.set.labels())
        return Subset(dataset, indices=list(range(self.__len__()//10*9, self.__len__())))


#ts = TimeseriesData()
#print(ts.values())
#print(ts.labels())

#sv = SupervisedTimeseriesData(lag=10)
#print(sv.values().shape)
#print(sv.labels().shape)

#svd = SupervisedTimeseriesDataset()
#print(len(svd.train_set()), len(svd.val_set()), len(svd.test_set()))
#print(svd.val_set())

#train_loader = DataLoader(svd.train_set(), batch_size=32, shuffle=True)
#valid_loader = DataLoader(dataset=svd.val_set(),batch_size=32,shuffle=False)
#test_loader = DataLoader(dataset=svd.test_set(),batch_size=32,shuffle=False)
